import abc
from collections.abc import Iterable, Iterator
from typing import Generic, TypeVar

from zkl_aiutils_training import Dock, ProcessingTask

from zkl_ptutils_training.plugins import MetricOrValue, MetricsCollector

AnyInput = TypeVar('AnyInput')
AnyOutput = TypeVar('AnyOutput')


class MLProducer(Generic[AnyInput], abc.ABC):
    @abc.abstractmethod
    def produce(self) -> AnyInput:
        pass


class MLProcessor(Generic[AnyInput, AnyOutput], abc.ABC):
    @abc.abstractmethod
    def process(self, input: AnyInput) -> AnyOutput:
        pass


class MLProcessingTask(ProcessingTask[AnyInput, AnyOutput]):
    def __init__(self, *,
        producer: MLProducer[AnyInput] | Iterable[AnyInput] | Iterator[AnyInput],
        processor: MLProcessor[AnyInput, AnyOutput],
    ):
        super().__init__()
        if not isinstance(producer, MLProducer):
            from .producer_simple import SimpleProducer
            producer = SimpleProducer(producer)
        self._producer = producer
        self._processor = processor

    def on_installed(self, dock: Dock):
        super().on_installed(dock)
        self.dock.install(self._producer)
        self.dock.install(self._processor)

    def _produce(self) -> AnyInput:
        return self._producer.produce()

    def _process(self, input: AnyInput) -> AnyOutput:
        return self._processor.process(input)

    def _consume(self, output: AnyOutput):
        pass  # do nothing


class MLValidatingTask(MLProcessingTask[AnyInput, AnyOutput], MetricsCollector):
    def run(self) -> dict[str, MetricOrValue]:
        super().run()
        return self._get_metrics()
